#include "greedy-allocation.h"
#include "ns3/log.h"
#include <iomanip>  // 添加此头文件以支持 std::setprecision
#include <algorithm>
#include <vector>
#include <utility>

namespace ns3 {

NS_LOG_COMPONENT_DEFINE("GreedyAllocation");

NS_OBJECT_ENSURE_REGISTERED(GreedyAllocation);

TypeId GreedyAllocation::GetTypeId()
{
    static TypeId tid = TypeId("ns3::GreedyAllocation")
        .SetParent<ChannelAllocation>()
        .SetGroupName("SAGINIoV")
        .AddConstructor<GreedyAllocation>();
    return tid;
}

GreedyAllocation::GreedyAllocation()
{
    NS_LOG_FUNCTION(this);
}

GreedyAllocation::~GreedyAllocation()
{
    NS_LOG_FUNCTION(this);
}

std::map<uint32_t, int> GreedyAllocation::AllocateChannels()
{
    NS_LOG_FUNCTION(this);
    
    std::map<uint32_t, int> allocation;
    
    if (!m_nodeManager) {
        NS_LOG_ERROR("Node manager is not set");
        return allocation;
    }
    
    // 获取当前仿真时间
    double currentTime = Simulator::Now().GetSeconds();
    
    // 获取当前活跃的车辆数据
    std::vector<VehicleData> activeVehicles = m_nodeManager->GetVehicleDataAtTime(currentTime);
    NS_LOG_INFO("Number of active vehicles at time " << currentTime << ": " << activeVehicles.size());
    
    // 获取所有车辆节点
    NodeContainer vehicles = m_nodeManager->GetVehicleNodes();
    
    // 确定可用信道的数量和范围
    NodeContainer groundStations = m_nodeManager->GetGroundStationNodes();
    NodeContainer satellites = m_nodeManager->GetSatelliteNodes();
    
    int numChannels = 0;
    if (m_allowGroundStationConnection) {
        numChannels += groundStations.GetN();
    }
    if (m_allowSatelliteConnection) {
        numChannels += satellites.GetN();
    }
    
    if (numChannels == 0) {
        NS_LOG_WARN("No available channels");
        return allocation;
    }
    
    // 为每个活跃车辆找到最佳信道
    for (const auto& vehicle : activeVehicles) {
        // 查找对应的节点ID
        uint32_t nodeId = 0;
        for (uint32_t i = 0; i < vehicles.GetN(); i++) {
            std::string vehicleId = m_nodeManager->GetVehicleIdForNode(vehicles.Get(i)->GetId());
            if (vehicleId == vehicle.id) {
                nodeId = vehicles.Get(i)->GetId();
                break;
            }
        }
        
        if (nodeId == 0) {
            NS_LOG_WARN("Could not find node for vehicle " << vehicle.id);
            continue;
        }
        
        // 尝试所有可能的信道，找到得分最高的
        double bestScore = -std::numeric_limits<double>::infinity();
        int bestChannelId = -1;
        
        for (int channelId = 0; channelId < numChannels; channelId++) {
            // 检查信道类型是否允许
            bool isGroundStationChannel = channelId < static_cast<int>(groundStations.GetN());
            if ((isGroundStationChannel && !m_allowGroundStationConnection) ||
                (!isGroundStationChannel && !m_allowSatelliteConnection)) {
                continue;
            }
            
            double score = CalculateChannelScore(nodeId, channelId, allocation);
            if (score > bestScore) {
                bestScore = score;
                bestChannelId = channelId;
            }
        }
        
        if (bestChannelId >= 0) {
            allocation[nodeId] = bestChannelId;
            NS_LOG_DEBUG("Vehicle " << vehicle.id << " (Node " << nodeId << ") allocated to channel " << bestChannelId << " with score " << bestScore);
        }
    }
    
    return allocation;
}

double GreedyAllocation::CalculateChannelScore(uint32_t nodeId, int channelId, const std::map<uint32_t, int>& currentAllocation)
{
    NS_LOG_FUNCTION(this << nodeId << channelId);
    
    // 临时添加当前分配，用于计算SINR
    std::map<uint32_t, int> tempAllocation = currentAllocation;
    tempAllocation[nodeId] = channelId;
    m_currentAllocation = tempAllocation; // 临时设置，用于SINR计算
    
    // 计算SINR
    double sinr = CalculateSINR(nodeId, channelId);
    
    // 计算负载因子
    double loadFactor = 0.0;
    
    // 确定连接的节点类型
    NodeContainer groundStations = m_nodeManager->GetGroundStationNodes();
    bool isGroundStationChannel = channelId < static_cast<int>(groundStations.GetN());
    
    if (isGroundStationChannel) {
        // 获取地面站节点ID
        uint32_t stationNodeId = groundStations.Get(channelId)->GetId();
        loadFactor = m_nodeManager->GetGroundStationLoad(stationNodeId);
    } else {
        // 获取卫星节点ID
        NodeContainer satellites = m_nodeManager->GetSatelliteNodes();
        int satelliteIndex = channelId - groundStations.GetN();
        if (satelliteIndex < static_cast<int>(satellites.GetN())) {
            uint32_t satelliteNodeId = satellites.Get(satelliteIndex)->GetId();
            loadFactor = m_nodeManager->GetSatelliteLoad(satelliteNodeId);
        }
    }
    
    // 计算综合得分: SINR - 负载权重 * 负载因子
    double loadWeight = 10.0; // 负载权重，调整SINR和负载的相对重要性
    double score = sinr - loadWeight * loadFactor;
    
    // 恢复原先的分配
    m_currentAllocation = currentAllocation;
    
    return score;
}

// 在文件末尾添加 LogAllocationResults 的实现
void GreedyAllocation::LogAllocationResults()
{
    NS_LOG_FUNCTION(this);
    
    Time now = Simulator::Now();
    std::cout << "\n======== Channel Allocation at t=" << now.GetSeconds() << "s ========" << std::endl;
    std::cout << "Algorithm: Greedy Allocation" << std::endl;
    
    // 打印每个车辆的分配信息
    NodeContainer vehicles = m_nodeManager->GetVehicleNodes();
    NodeContainer groundStations = m_nodeManager->GetGroundStationNodes();
    NodeContainer satellites = m_nodeManager->GetSatelliteNodes();
    
    for (auto it = m_currentAllocation.begin(); it != m_currentAllocation.end(); ++it) {
        uint32_t nodeId = it->first;
        int channelId = it->second;
        
        // 获取车辆ID
        std::string vehicleId = m_nodeManager->GetVehicleIdForNode(nodeId);
        
        // 确定连接的是地面站还是卫星
        std::string stationType = "Unknown";
        uint32_t stationId = 0;
        
        if (channelId < static_cast<int>(groundStations.GetN())) {
            stationType = "Ground Station";
            stationId = groundStations.Get(channelId)->GetId();
        } else {
            stationType = "Satellite";
            int satelliteIndex = channelId - groundStations.GetN();
            if (satelliteIndex < static_cast<int>(satellites.GetN())) {
                stationId = satellites.Get(satelliteIndex)->GetId();
            }
        }
        
        // 计算当前SINR
        double sinr = CalculateSINR(nodeId, channelId);
        
        std::cout << "Vehicle " << vehicleId << " (Node " << nodeId << "): Channel " << channelId 
                  << " (Connected to " << stationType << " " << stationId << "), SINR: " 
                  << std::fixed << std::setprecision(2) << sinr << " dB" << std::endl;
    }
    
    std::cout << "Total vehicles allocated: " << m_currentAllocation.size() << std::endl;
    std::cout << "=======================================================" << std::endl;
}

} // namespace ns3
